<script setup lang="ts">
import { onMounted, ref } from "vue";
import { createLabelsDataset, showOCRImages } from "./OCRDataset";
import { saveDir, imagesFileName } from "./OCRDatasetConstants";
import { ElButton, ElTree, ElIcon, ElProgress } from "element-plus";
import { Files, Folder } from "@element-plus/icons-vue";
import { getDir, getFile, createFile } from "@/renderer/api/messageAPI";
import { dir2tree } from "@/renderer/utils/utils";
import createImagesDatasetURL from "./createImagesDataset.worker?worker&url";
import Worker from "@/renderer/utils/Worker";

const imgBox = ref();
const datasetsTree = ref([]);
const createImagesDatasetProgress = ref(0)

const makeTree = () => {
  getDir("datasets").then((res) => {
    datasetsTree.value = dir2tree(res);
  });
};

const handleNodeClick = (data, node) => {
  if (!node.isLeaf) return;

  imgBox.value.innerHTML = "";

  getFile("datasets/" + data.path).then((res) => {
    if (res.byteLength && /ocr.+?images/.test(data.path)) {
      imgBox.value.append(showOCRImages(res));
    }
  });
};

const makeOCRDataset = async () => {
  imgBox.value.innerHTML = "";

  const worker = new Worker(createImagesDatasetURL);
  worker.onMessage((event) => {
    if (event.data.type == 'createImagesDatasetProgress') createImagesDatasetProgress.value = event.data.data
  })
  const imagesDataset = await worker.postMessage({
    type: "createImagesDataset",
  });
  worker.terminate()
  await createFile(saveDir + imagesFileName, imagesDataset);

  await createLabelsDataset();

  makeTree();
};

onMounted(() => {
  makeTree();
});
</script>

<template>
  <div class="container">
    <ElButton plain round type="primary" @click="makeOCRDataset">
      生成OCR数据集
    </ElButton>
    <ElTree class="tree" :data="datasetsTree" @node-click="handleNodeClick">
      <template #default="{ node, data }">
        <ElIcon v-if="data.children.length > 1" size="15"><Folder /></ElIcon>
        <ElIcon v-else size="15"><Files /></ElIcon>
        <span>{{ node.label }}</span>
      </template>
    </ElTree>
    <ElProgress
      style="width: 350px;"
      :percentage="createImagesDatasetProgress"
      :stroke-width="20"
      :text-inside="true"
      striped
      :striped-flow="createImagesDatasetProgress < 100"
    />
    <div ref="imgBox"></div>
  </div>
</template>

<style scoped>
.container {
  display: flex;
  flex-direction: column;
  align-items: start;
}

.container > * {
  margin: 10px;
}

.tree {
  min-width: 300px;
  border-radius: 5px;
  padding: 10px;
}

.el-icon {
  margin: 0 5px;
}
</style>
